#include "CudaEpiModel.h"
#include "CudaParam.h"
#include <thrust/count.h>

void cudaResponse(){
	thrust::for_each(tract_vec->begin(),tract_vec->end(),ActOp0(model->nTimer));
	
	err = cudaThreadSynchronize();
        if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;

	// 1. Count cumulative number of ascertained cases
	if(!model->bTrigger){
		thrust::device_vector<int> comm_nNumAscertained_vec(model->nNumCommunities);
		thrust::transform(comm_vec->begin(),comm_vec->end(),comm_nNumAscertained_vec.begin(),ActOp1());
		int nNumAscertained = thrust::reduce(comm_nNumAscertained_vec.begin(),comm_nNumAscertained_vec.end());
	
		if (nNumAscertained/(double)model->nNumPerson>model->fResponseThreshhold) {
			model->bTrigger=true;
			model->nTriggerTime=model->nTimer+model->nTriggerDelay*2;
			cout << "Response starts on day " << (model->nTriggerTime/2) << ": ascertained: " 
			<< nNumAscertained << "/" << model->nNumPerson << "=" << nNumAscertained/(double)model->nNumPerson 
			<< " on day " << (model->nTimer/2) << endl;
		}
	}
	
	err = cudaThreadSynchronize();
        if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	// 2. Epidemic is detected, start reactive vaccination strategies
	if(model->bTrigger && model->nTimer==model->nTriggerTime) {

		thrust::for_each(tract_vec->begin(),tract_vec->end(),
			ActOp2(thrust::raw_pointer_cast(ppl_vec->data()),
				thrust::raw_pointer_cast(comm_vec->data()),
				nNumWantVaccine,
				model->eAVPolicy,model->fQuarantineCompliance,model->nSchoolClosureDays,
				model->nSchoolClosurePolicy,model->eVaccinationStrategy,model->nAscertainmentDelay));
			
	}
	
	err = cudaThreadSynchronize();
        if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	 // 3. change vaccine priorities if needed
	if (model->nPriorityChangeTime>0 && model->nPriorityChangeTime==model->nTimer/2) {
		thrust::for_each(ppl_vec->begin(),ppl_vec->end(),
			ActOp3(thrust::raw_pointer_cast(ppl_vec->data()),
				   model->nNumPerson,model->nHighestPriority
				   ));
		
	}
	
	err = cudaThreadSynchronize();
        if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	// 4. Ring vaccination and consider closing or re-opening schools
	if (model->bTrigger && model->nTimer>=model->nTriggerTime) {
		if (model->eVaccinationStrategy==RINGVACTRACT) {
	
			thrust::for_each(tract_vec->begin(),tract_vec->end(),
				ActOp41(thrust::raw_pointer_cast(ppl_vec->data()),
						thrust::raw_pointer_cast(comm_vec->data()),
						nNumWantVaccine,
						model->nAscertainmentDelay));
						
		}
		
		if (model->nSchoolClosureDays>0) {
			if (model->nSchoolClosurePolicy==2) { // close schools by ascertainment
				thrust::for_each(tract_vec->begin(),tract_vec->end(),
					ActOp42(thrust::raw_pointer_cast(ppl_vec->data()),
							thrust::raw_pointer_cast(comm_vec->data()),
							model->nAscertainmentDelay));
				// original program has output string
			}
			
			thrust::for_each(tract_vec->begin(),tract_vec->end(),ActOp43(model->nTimer));
			// original program has output string
		}
	}
	
	err = cudaThreadSynchronize();
        if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	// 5. Count available vaccines
	for (int i=0; i<model->nNumVaccineTypes; i++)
		if (model->vaccineproductionschedule[i][model->nTimer/2]>0)
			model->nVaccineSupply[i] += model->vaccineproductionschedule[i][model->nTimer/2];  // new vaccine arrives

	unsigned int nVacLim[NUMVACCINES];
	unsigned int totalvaccinesupply=0;
	
	for (int i=0; i<model->nNumVaccineTypes; i++) {
		nVacLim[i] = model->nVaccineSupply[i]-model->nNumVaccineDosesUsed[i]; // how many vaccine doses are available today?
		if (model->nVaccineSupply[i]<=model->nNumVaccineDosesUsed[i]) // because we are using unsigned ints
			nVacLim[i]=0;
		totalvaccinesupply += nVacLim[i];
	}
	if (totalvaccinesupply>model->nVaccineDailyLimit) { // limited overall vaccine daily delivery capacity
		double factor = model->nVaccineDailyLimit/(double)totalvaccinesupply;
		for (int i=0; i<model->nNumVaccineTypes; i++) nVacLim[i] *= factor;
	}
	
	err = cudaThreadSynchronize();
        if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	// 6. Distribute vaccines
	thrust::device_vector<unsigned int> nVacLim_vec(model->nNumPerson);
	thrust::device_vector<unsigned int> d_nNumVaccineDosesUsed(
		model->nNumVaccineDosesUsed,model->nNumVaccineDosesUsed+NUMVACCINES);
	
	
	for (int vacnum=0; vacnum<model->nNumVaccineTypes; vacnum++) {
		
		if (nVacLim[vacnum]>0) {
			double vFrac=0.0;
			
			thrust::device_vector<unsigned int> wantthisvac(PRIORITY_LAST,0);
			thrust::for_each(ppl_vec->begin(),ppl_vec->end(),
				ActOp61(thrust::raw_pointer_cast(wantthisvac.data()),vacnum));
				
			for (int priority=1; priority<=model->nHighestPriority; priority++) {
			
				unsigned int wantvac = wantthisvac[priority];
				if (wantvac > 0) {
					vFrac = nVacLim[vacnum]/(double)wantvac;
				}
				thrust::device_vector<int> mutex(1,0);
				thrust::for_each(ppl_vec->begin(),ppl_vec->end(),
					ActOp62(priority,vacnum,vFrac,
							thrust::raw_pointer_cast(rand_vec->data()),
							thrust::raw_pointer_cast(d_nNumVaccineDosesUsed.data()),
							thrust::raw_pointer_cast(nVacLim_vec.data()),
							nNumWantVaccine,
							model->ePrevaccinationStrategy,
							thrust::raw_pointer_cast(mutex.data())
							));
				
				unsigned int nVacLim_sum = thrust::reduce(nVacLim_vec.begin(),nVacLim_vec.end());
				if(nVacLim[vacnum] > nVacLim_sum) nVacLim[vacnum] -= nVacLim_sum;
				else nVacLim[vacnum] = 0;
				
				err = cudaThreadSynchronize();
					if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
			}
		}
	}
	
	err = cudaThreadSynchronize();
        if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	// 7. Count available antiviral courses
	cudaMemcpy(&model->nNumAntiviralsUsed,nNumAntiviralsUsed,sizeof(unsigned int),cudaMemcpyHostToDevice);
	cudaMemcpy(&model->nNumWantAV,nNumWantAV,sizeof(unsigned int),cudaMemcpyDeviceToHost);
	
	unsigned int uAVFraction;
	unsigned int nAVLim = min(model->nAVTotalLimit-model->nNumAntiviralsUsed, model->nAVDailyLimit);
	if (model->nAVTotalLimit <= model->nNumAntiviralsUsed) nAVLim=0;
	
	
	if (model->nNumWantAV <= nAVLim)
		uAVFraction = UINT_MAX;
	else
		uAVFraction = (unsigned int)(nAVLim/(double)model->nNumWantAV*UINT_MAX);
	
	thrust::for_each(ppl_vec->begin(),ppl_vec->end(),
		ActOp71(nNumAntiviralsUsed,nNumWantAV,uAVFraction,
			thrust::raw_pointer_cast(rand_vec->data())));
	
	cudaMemcpy(&model->nNumAntiviralsUsed,nNumAntiviralsUsed,sizeof(unsigned int),cudaMemcpyDeviceToHost);
	
	for (int i=0; i<model->nNumVaccineTypes; i++)
		model->nNumVaccineDosesUsedHistory[i].push_back(d_nNumVaccineDosesUsed[i]);
	model->nNumAntiviralsUsedHistory.push_back(model->nNumAntiviralsUsed);
	
	cout << "response finish" << endl;
	err = cudaThreadSynchronize();
        if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
}

void cudaStats(){
	unsigned int nNumSymptomatic=0;
    unsigned int nNumCumulativeSymptomatic=0;
	thrust::device_vector<unsigned int> results(model->nNumCommunities);
	
	thrust::transform(comm_vec->begin(),comm_vec->end(),results.begin(),StatsOp1());
	nNumSymptomatic = reduce(results.begin(),results.end());
	
	thrust::transform(comm_vec->begin(),comm_vec->end(),results.begin(),StatsOp2());
	nNumCumulativeSymptomatic = reduce(results.begin(),results.end());
	
	model->nNumCumulativeSymptomaticHistory.push_back(nNumCumulativeSymptomatic);
    model->nNumSymptomaticHistory.push_back(nNumSymptomatic);
	
	cout << "nNumSymptomatic: " << nNumSymptomatic << endl;
	cout << "nNumCumulativeSymptomatic: " << nNumCumulativeSymptomatic << endl;
}